﻿#pragma once

#include "worklet/EquationSet.h"

#include "Application.h"
#include "EulerEquation.h"
#include "Mesh.h"
#include "Transient.h"

#include "TopoTables.h"
#include <vtkm/Matrix.h>
#include <vtkm/VectorAnalysis.h>
#include <vtkm/cont/ArrayHandleGroupVecVariable.h>
#include <vtkm/worklet/WorkletCellNeighborhood.h>
#include <vtkm/worklet/WorkletMapField.h>
#include <vtkm/worklet/WorkletMapTopology.h>

class EulerEquation;
namespace Loops
{

struct FiniteVolumeWorklet : vtkm::worklet::WorkletVisitCellsWithPoints
{
  VTKM_CONT
  FiniteVolumeWorklet(FloatDefault t)
    : _t(t)
  {
  }

  using ControlSignature = void(CellSetIn cell_set_in,
                                FieldInPoint coords,
                                FieldInCell cell_link,
                                WholeArrayIn sln,
                                FieldOutCell rhs,
                                ExecObject context);

  using ExecutionSignature = void(CellShape, InputIndex, _2, _3, _4, _5, _6);

  template<typename PointCoordVecType>
  VTKM_EXEC Vec3f ComputeDirectedArea(const vtkm::CellShapeTagGeneric& shape,
                                      const PointCoordVecType& points) const // 有向面积
  {
    Vec3f directed_area{ 1, 1, 1 };
    switch (shape.Id)
    {
      case vtkm::CELL_SHAPE_VERTEX:
      {
        this->RaiseError("1D vertex normal unsolved.");
      }
      break;
      case vtkm::CELL_SHAPE_LINE:
      {
        using CellShape = vtkm::CellShapeIdToTag<vtkm::CELL_SHAPE_LINE>::Tag;
        auto _u_axis = points[1] - points[0];
        directed_area = { _u_axis[1], -_u_axis[0], 0 };
        //directed_area = vtkm::Cross(points[1] - points[0], Vec3f{ 0,0,1 });
      }
      break;

      case vtkm::CELL_SHAPE_TRIANGLE:
      {
        directed_area = vtkm::Cross(points[1] - points[0], points[2] - points[0]);
        directed_area = directed_area / 2;
      }
      break;
      case vtkm::CELL_SHAPE_QUAD:
      {
        directed_area = vtkm::Cross(points[0] - points[1], points[2] - points[1]);
      }
      break;
      default:
        this->RaiseError("unknown cell type %d");
        break;
    }
    return directed_area;
  }

  template<typename PointCoordVecType>
  VTKM_EXEC Vec3f FaceCenter(const vtkm::CellShapeTagGeneric& shape,
                             const PointCoordVecType& pts) const // 有向面积
  {
    Vec3f centroid{ 1, 1, 1 };
    switch (shape.Id)
    {
      case vtkm::CELL_SHAPE_VERTEX:
      {
        this->RaiseError("1D vertex normal unsolved.");
      }
      break;
      case vtkm::CELL_SHAPE_LINE:
      {
        centroid = (pts[0] + pts[1]) / 2;
      }
      break;

      case vtkm::CELL_SHAPE_TRIANGLE:
      {
        centroid = (pts[0] + pts[1] + pts[2]) / 3;
      }
      break;
      case vtkm::CELL_SHAPE_QUAD:
      {
        centroid = (pts[0] + pts[1] + pts[2] + pts[3]) / 4;
      }
      break;
      default:
        this->RaiseError("unknown cell type %d");
        break;
    }
    return centroid;
  }


  template<typename CellShapeTag,
           typename CoordType,
           typename CellLink,
           typename SlnType,
           typename RhsType,
           typename ContextType>
  VTKM_EXEC void operator()(CellShapeTag& cell_shape,
                            Id& cell_id,
                            const CoordType& pts,
                            const CellLink& cell_link,
                            const SlnType& sln,
                            RhsType& rhs,
                            ContextType& context) const
  {
    auto num_faces = TopoTables{}.NumSides(cell_shape.Id);
    Real volume = 0;
    rhs = vtkm::TypeTraits<RhsType>::ZeroInitialization();
    using ValueType = RhsType;
    ValueType ul, ur;
    RhsType flux;

    Vec3f f2n[4];
    for (IdComponent side = 0; side < num_faces; side++)
    {
      UInt8 face_shape;
      CellFaceShape(side, cell_shape, face_shape);
      IdComponent n_f2n;
      CellFaceNumberOfPoints(side, cell_shape, n_f2n);

      for (IdComponent i = 0; i < n_f2n; i++)
      {
        IdComponent result;
        CellFaceLocalIndex(i, side, cell_shape, result);
        f2n[i] = pts[result];
      }

      ValueType neighbor_value;
      Id neighbor_cell_id = cell_link[side].Get().first;
      Vec3f centroid = FaceCenter(face_shape, f2n);

      Vec3f normal = ComputeDirectedArea(face_shape, f2n);
      const Real face_area = vtkm::Magnitude(normal);
      normal = normal / face_area;

      ul = sln.Get(cell_id);
      ur = sln.Get(cell_id);
      if (neighbor_cell_id < 0) // 边界条件
      {
        auto sid = cell_link[side].Get().second;
        BCType bc_type = static_cast<BCType>(sid);
        ur = context.BoundaryCondition(bc_type, ul, normal, centroid);
      }
      else
      {
        ur = sln.Get(neighbor_cell_id);
      }
      context.RiemannFlux(ul, ur, normal, flux);
      rhs += face_area * flux;
      volume += normal[1] * centroid[1] * face_area;
    }
    //std::cout << volume << std::endl;
    //printf("volume: %f \n", volume);

    rhs = rhs / (volume);
  }

  Real _t;
};

void ComputeResidual(EulerEquation& eq);
void ComputeResidualIsoVortex(EulerEquation& eq);
void ComputeResidualIsoVortexOnUnstructuredMesh(EulerEquation& eq);

void SolutionToArray(EulerEquation& eq);
void ExplictFirtOrder(EulerEquation& eq);
}